"""
reAct 大模型的一种思维模式，让大模型先思考再行动。（Reason + Act） 框架。
主要工具：
    tavily 搜索检索 ，需要配置api
    langsmit  agent 监控工具 ，需要配置api

本示例：
    使用 LangGraph 构建一个完整的“ReAct智能体工作流”。
    - Tavily 搜索工具：让模型具备网络搜索能力。
    - LangSmith（可选）：用于可视化 agent 的执行过程。

"""
import asyncio
import json
import operator
from typing import TypedDict, List, Annotated, Tuple, Union, Literal

from dotenv import load_dotenv
from langchain import hub
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from pydantic import Field, BaseModel, ValidationError

# 自动加载 .env 文件中的环境变量
load_dotenv()

# ---------------- 1️⃣ 定义工具 ----------------
tools = [TavilySearchResults(max_results=1)]  # 创建 Tavily 搜索工具（可联网搜索），每次只返回一个搜索结果

# ---------------- 2️⃣ 拉取 prompt 模板 ----------------
prompt = hub.pull("wfh/react-agent-executor")  # 从 LangChain Hub 获取 ReAct 框架提示模板  ,Hub 上保存了一些通用的 prompt，可以直接 pull 使用
prompt.pretty_print()

# ---------------- 3️⃣ 定义大模型 ----------------
model_name = "xdeepseekv32exp"  # 讯飞星火 "xdeepseekv32exp"

# 定义大语言模型
llm = ChatOpenAI(model=model_name)

# ---------------- 4️⃣ 创建 ReAct agent ----------------
agent_executor = create_react_agent(llm, tools, prompt=prompt)  # 会自动管理“思考 → 行动 → 观察 → 再思考”的循环逻辑


# ---------------- 5️⃣ 定义状态结构 ----------------
class PlanExecute(TypedDict):
    """状态字典，用于存储输入、计划、已完成步骤和响应"""
    input: str  # 用户录入提示词
    plan: List[str]  # 用户录入提示词如果是大问题，会被拆分成几个小问题
    past_steps: Annotated[List[Tuple], operator.add]  # 已执行的步骤及结果
    response: str  # 最终回复给用户的答案


# ---------------- 6️⃣ 定义 Pydantic 数据模型 ----------------
## 计划模型
class Plan(BaseModel):
    """表示一组待执行的计划步骤"""
    steps: List[str] = Field(description="需要执行的不同步骤，应该按顺序排列")


## 输出模型
class Response(BaseModel):
    """直接回答"""
    response: str


##  要实行的操作
class Act(BaseModel):
    """要执行的动作"""
    action: Union[Response, Plan] = Field(
        description="若已可回答则返回 Response，否则返回新的 Plan。"
    )


# ---------------- 7️⃣ 定义计划生成器 ----------------
planner_prompt = ChatPromptTemplate.from_messages([
    ("system",
     """
     你是一个任务规划助手。请将用户的需求分解为具体的执行步骤。

     要求返回严格 JSON 格式：
     {
       "steps": ["步骤1", "步骤2", "步骤3"]
     }

     注意：
     - 每个步骤应该是独立、可执行的任务
     - 不要添加多余步骤
     - 最后一步应能直接得到最终答案
     - 只返回 JSON，不要添加任何其他文字
     """),
    ("human", "用户输入: {input}")
])

async def plan_step(state: PlanExecute):
    """计划步骤 - 手动处理 JSON 解析"""
    try:
        # 调用模型生成计划
        response = await llm.ainvoke(planner_prompt.format_messages(input=state["input"]))
        raw_output = response.content.strip()
        print(f"🔍 Planner 原始输出:\n{raw_output}")

        # 清理输出，提取 JSON 部分
        if "```json" in raw_output:
            # 提取代码块中的 JSON
            json_str = raw_output.split("```json")[1].split("```")[0].strip()
        elif "```" in raw_output:
            # 提取其他代码块
            json_str = raw_output.split("```")[1].split("```")[0].strip()
        else:
            # 直接使用内容
            json_str = raw_output

        # 解析 JSON
        plan_data = json.loads(json_str)
        steps = plan_data.get("steps", [])

        print(f"📋 解析后的步骤: {steps}")
        return {"plan": steps}

    except Exception as e:
        print(f"❌ 计划生成失败: {e}")
        # 返回默认计划
        return {"plan": [f"搜索关于'{state['input']}'的信息"]}


# ---------------- 8️⃣ 定义执行步骤 ----------------
async def execute_step(state: PlanExecute):
    """执行计划中的当前步骤"""
    if not state.get("plan"):
        return {"past_steps": state["past_steps"]}

    plan = state["plan"]
    plan_str = "\n".join(f"{i + 1}. {step}" for i, step in enumerate(plan))
    current_task = plan[0]  # 执行第一个任务

    task_formatted = f"""以下是完整的执行计划：
{plan_str}

请执行第 1 步任务：{current_task}"""

    print(f"🎯 执行任务: {current_task}")

    try:
        # 调用 ReAct agent 执行任务
        agent_response = await agent_executor.ainvoke({"messages": [("user", task_formatted)]})
        result_content = agent_response["messages"][-1].content

        print(f"✅ 任务完成结果: {result_content[:100]}...")  # 只打印前100字符

        return {
            "past_steps": state["past_steps"] + [(current_task, result_content)]
        }
    except Exception as e:
        print(f"❌ 任务执行失败: {e}")
        return {
            "past_steps": state["past_steps"] + [(current_task, f"执行失败: {str(e)}")]
        }


# ---------------- 9️⃣ 定义重新计划器 ----------------
replanner_prompt = ChatPromptTemplate.from_template("""
你是任务规划助手。根据用户目标、当前计划和已完成步骤，生成下一步行动。

要求返回严格 JSON，格式如下：
{
  "action": {
    "response": "string，可选。如果任务完成，直接返回最终回答。",
    "steps": ["string", "string", ...] 可选。如果还需要执行，返回剩余计划步骤。
  }
}

⚠️ 注意：
- 若任务完成，返回 {"action": {"response": "最终回答"}}
- 若任务未完成，返回 {"action": {"steps": ["剩余任务1", "任务2"]}}
- 不要输出除 JSON 外的任何文字

用户输入：
{input}

原计划：
{plan}

已完成步骤：
{past_steps}
""")


# 自行规划重新执行的函数
async def replan_step(state: PlanExecute):
    """重新规划任务"""
    try:
        # 格式化提示词
        messages = replanner_prompt.format_messages(
            input=state["input"],
            plan=state.get("plan", []),
            past_steps=state.get("past_steps", [])
        )

        # 调用模型
        response = await llm.ainvoke(messages)
        raw_output = response.content.strip()
        print(f"🔍 Replanner 原始输出:\n{raw_output}")

        # 清理和解析 JSON
        if "```json" in raw_output:
            json_str = raw_output.split("```json")[1].split("```")[0].strip()
        elif "```" in raw_output:
            json_str = raw_output.split("```")[1].split("```")[0].strip()
        else:
            json_str = raw_output

        data = json.loads(json_str)

        # 根据返回内容决定下一步
        action_data = data.get("action", {})

        if "response" in action_data and action_data["response"]:
            # 有最终答案，结束流程
            return {"response": action_data["response"]}
        elif "steps" in action_data and action_data["steps"]:
            # 有新的步骤，继续执行
            return {"plan": action_data["steps"]}
        else:
            # 默认情况，结束流程
            return {"response": "无法确定下一步行动"}

    except Exception as e:
        print(f"❌ 重新计划失败: {e}")
        return {"response": f"规划过程出现错误: {str(e)}"}


# ---------------- 🔟 定义条件判断 ----------------
def should_end(state: PlanExecute):
    """判断工作流是否应该结束"""
    if state.get("response"):
        return "__end__"  # 有最终响应，结束
    else:
        return "agent"  # 继续执行


# ---------------- 11️⃣ 主函数 ----------------
async def main():
    # 创建状态图
    workflow = StateGraph(PlanExecute)

    # 添加节点
    workflow.add_node("planner", plan_step)
    workflow.add_node("agent", execute_step)
    workflow.add_node("replan", replan_step)

    # 添加边
    workflow.add_edge(START, "planner")
    workflow.add_edge("planner", "agent")
    workflow.add_edge("agent", "replan")
    workflow.add_conditional_edges("replan", should_end)

    # 编译应用
    app = workflow.compile()

    # 保存流程图
    try:
        graph_png = app.get_graph().draw_mermaid_png()
        with open("llm_langgraph_1_agent工作流.png", "wb") as f:
            f.write(graph_png)
        print("📊 流程图已保存")
    except Exception as e:
        print(f"⚠️ 流程图保存失败: {e}")

    # 执行工作流
    inputs = {"input": "2024年巴黎奥运会100米自由泳决赛冠军的家乡是哪里？请用中文答复"}
    config = {"recursion_limit": 50}

    print("🚀 开始执行工作流...")
    async for event in app.astream(inputs, config=config):
        for k, v in event.items():
            if k != "__end__":
                print(f"📦 状态更新 - {k}: {v}")

    print("🎉 工作流执行完成!")


# ---------------- 12️⃣ 运行 ----------------
if __name__ == "__main__":
    asyncio.run(main())